#ifndef _KeyedLists_H_
#define _KeyedLists_H_

//
// KeyedLists.h
//
// Created by TB on 11/17/11.
// Copyright 2011 BNL. All rights reserved.
//

#include <map>
#include <utility>

#include <TH1D.h>

namespace tppmc {
   
   
   /**
    \todo This is really the wrong way around - each PARTON should have a
    hadron/Q list. That way the KeyedLists itself is normalised internally,
    rather than being cross-normalised with other KeyedLists.
    Simple wrapper around std::map<int, std::vector<double> >.
    The key (int) represents a parton species.
    The value (vector<double>) represents a collection of Q (NOT Q^2) bins.
    It is used by fragmentation functions to store the probabilities for a
    hadron to be produced from a particular parton at a particular Q.
    Each hadron species stores a 
    A collection of lists, keyed by parton type.
    Each list represents a collection of Q bins.
    */
   class KeyedLists {
      
   public:
      
      
      /**
       Default constructor.
       A collection of bin lists, each keyed by an integer.
       Initialise with the number of bins and their width.
       The bins are of uniform in width from min to max.
       Each entry in the table has the same bins.
       */
      KeyedLists(int nBins = 50,
                 double min = 2.,
                 double max = 102.);
      
      
      /**
       Destructor.
       */
      virtual ~KeyedLists() { }
      
      
      /**
       Add an entry to the table for a parton species.
       Input: parton ID code according to KKP convention.
       Returns true if the entry is added or false if an entry
       for this parton type already exists.
       */
      virtual bool Insert(int parton);
      
      
      /**
       Returns the probability stored for the specified parton
       and Q bin.
       Returns NaN if the parton and/or Q bin does not exist.
       */
      virtual double At(int parton, unsigned QBin) const;
      
      
      /**
       See At(int, int).
       */
      virtual double operator()(int parton, unsigned QBin) const;
      
      
      /**
       Returns true if an entry for the parton and Q bin exists, false
       if not.
       */
      virtual bool Exists(int parton, unsigned QBin) const;
      
      
      /**
       Set the value for a parton in a Q bin.
       The value is that at the lower edge of the bin. 
       */
      virtual void Set(int parton, unsigned QBin, double value);
      
      
      /**
       Return the probability for the requested parton at
       the requested Q value.
       Interpolation is done for values between bin edges.
       */
      virtual double Interpolate(int parton, double Q) const;
      
      
      /**
       Returns the (uniform) Q bin witth (GeV^2).
       */
      virtual double BinSize() const;
      
      
      virtual double Min() const { return mQMin; }
      virtual double Max() const { return mQMax; }
      virtual int NBins() const { return mNQBins; }
      /**
       */
#if 0
      /**
       Returns a vector of Q bin low-edges used in the table.
       */
      virtual std::vector<double> QBins() const;
      /**
       Returns a vector of parton types used in the table.
       */
      virtual std::vector<int> Partons() const;
#endif
   protected:

      unsigned mNQBins;
      double mQMin;
      double mQMax;
      
      typedef std::vector<double> Array;
      
      typedef std::map<int, Array> Table;
      
      Table mTable; ///< Store an array of probabilities
                    ///< (Q bins) for each parton type.
      
      ClassDef(tppmc::KeyedLists, 1)
   };
   
   inline double KeyedLists::BinSize() const {
      return (mQMax - mQMin) / mNQBins;
   }
   
   
   /**
    Wrapper around std::map<T, std::vector<U> > with some added functionality.
   */
   template<typename T, typename U>
   class Table {
   public:

      /**
       Initialise the table with a copy of vector<U> for each element in
       vector<T>
      */
      Table(const std::vector<T>&, const std::vector<U>&) {
      }
      
      /** Destructor */
      virtual ~Table() { }
      
      /**
       Add an entry for t, initialising each element of the vector<U> for
       that entry with the value u.
       Return true if the entry is added, false if not.
      */
      virtual bool Expand(const T&, const U&) { return false; }

   protected:
      std::map<T, std::vector<U> > mEntries;
   };


   /**
    (Potential replacement for KeyedLists, though Alan is working on
    removing cached integrals altogether.)
    Probabilities for a parton to fragment into
    different hadron species at different Q values.
    */
   class FragmentationTable : public TObject {
      
   public:
      
      /** Destructor */
      virtual ~FragmentationTable();
      
      /**
       Initialise the table with uniform binning in Q.
       Set nBins Q bins, each of width (max-min)/nBins.
       @param parton  Identifier for the table
       @param nBins   Number of Q-squared bins
       @param min     Minimum Q-squared
       @param max     Maximum Q-squared
       @param hadrons Optional list of hadron species
       */
      FragmentationTable(int nBins = 50,
                         double min = 2.,
                         double max = 102.,
                         const std::vector<int>& hadrons = std::vector<int>());
      
      /**
       \overload
       Initialise the table with non-uniform binning in Q.
       @param parton  Identifier for the table
       @param edges A list of Q bin edges. The number of resultant bins is
       edges.size() - 1. Each entry in edges must be larger than the last.
       */
      FragmentationTable(std::vector<double> edges,
                         const std::vector<int>& hadrons = std::vector<int>());

      /**
       Normalise the current table contents such that, in each Q bin,
       the sum over all hadrons equals norm.
       */
      virtual void Normalise(double norm = 1.);

      /**
       Returns the value for species in the bin containing Q.
       Returns NaN if there is no entry for that bin.
       */
      virtual double At(int species, double Q) const;

      /**
       Set the value for species in the bin containing Q.
       */
      virtual bool Set(int species, double Q, double value);

      /**
       Returns true if a value is stored for species and the bin
       containing Q.
       */
      virtual bool Contains(int species, double Q) const;

      /**
       Expand the table with an entry for species.
       Binning in Q is the same as for all other entries.
       Initialise each Q bin with the value initial.
       Returns true if the new entry is inserted, false if not.
       */
      virtual bool Insert(int species, double initial = 0.);

      /**
       Return a random hadron species at the requested Q.
       Should be normalised to 1 before calling.
       \todo Check whether interpolation assumes values are at bin edge or
       centre and make sure the calculation matches this.
       */
      virtual std::pair<int, bool> Random(double Q);

      /**
       Print the current contents of the table to standard output.
       Inherited from TObject.
       */
      virtual void Print(Option_t* = "") const;

   protected:

      /**
       Returns the index of the Q-squared bin containing Q.
       Indices are in the range [0, N).
       Returns -1 if Q does not fall in a bin.
       */
      virtual int Bin(double Q) const;

      // Logically, it might be nice to have a list of hadrons keyed by Q
      // bin, so the list values add up to total probability 1, but it's
      // easy to compute Q bin numerically and key by hadron type.
      // Use TH1D for as it has non-uniform binning and linear interpolation
      // already built in.
//      typedef std::vector<double> Array;
      typedef TH1D Array;

      mutable Array mPrototype; ///< Binning used for each entry

      typedef std::map<int, Array> Table;
      typedef Table::value_type Entry;

      Table mTable; ///< Lists of Q bins keyed by hadron species

      ClassDef(tppmc::FragmentationTable, 1)
   };
   
} // namespace tppmc

#endif
